﻿#pragma once

#include <learn/check.h>
#include <learn/math/math.h>
#include <learn/math/matrix.h>

#include <tuple>
#include <vector>

namespace xi
{

namespace math
{

/**
 * @brief 进退法搜索
 *
 * @tparam T
 * @param[in] f
 * @param[in] x 初始位置
 * @param[in] d 搜索方向
 * @param[in] h 步长
 * @return std::tuple<double, double>
 */
template <class T> std::tuple<double, double> inline forward_backward_line_search(double (*f)(T), T x, T d, double h)
{
    int t = 2;

    double a = 0;
    double b = h;
    double fx0 = f(x);

    int k = 0;
    while (k < 30)
    {
        double fx1 = f(x + d * b);
        if (fx1 < fx0)
            b *= t;
        else
        {
            if (k != 0)
                break;

            // 如果第一步上升，则记录上升点，反向搜索高点
            a = b;
            b = -h;
        }

        fx0 = fx1;
        k++;
    }

    return {std::min(a, b), std::max(a, b)};
}

/**
 * @brief 黄金分割法搜索
 *
 * @tparam T
 * @param[in] f
 * @param[in] x 初始位置
 * @param[in] d 搜索方向
 * @param[in] tol
 * @return double
 */
template <class T> inline double golden_section_line_search(double (*f)(T), T x, T d, double tol)
{
    static const double rb = (std::sqrt(5.0) - 1) / 2;
    static const double ra = 1 - rb;

    auto [a, b] = forward_backward_line_search<T>(f, x, d, 1e-2);

    // 选取试探点
    double lambda = a + ra * (b - a);
    double mu = a + rb * (b - a);

    // 计算试探点的值
    std::vector<double> vals = {a, lambda, mu, b};
    std::vector<double> fs = {f(x + d * a), f(x + d * lambda), f(x + d * mu), f(x + d * b)};

    // 同时比较 a, b, lambda, mu 处的值
    std::size_t t = std::distance(fs.begin(), std::min_element(fs.begin(), fs.end()));

    // 迭代
    while (vals[3] - vals[0] > tol)
    {
        if (t < 2)
        {
            vals[3] = vals[2];
            vals[2] = vals[1];
            vals[1] = vals[0] + ra * (vals[3] - vals[0]);
            fs[3] = fs[2];
            fs[2] = fs[1];
            fs[1] = f(x + d * vals[1]);

            if (fs[t + 1] < fs[t])
                t++;
        }
        else
        {
            vals[0] = vals[1];
            vals[1] = vals[2];
            vals[2] = vals[0] + rb * (vals[3] - vals[0]);
            fs[0] = fs[1];
            fs[1] = fs[2];
            fs[2] = f(x + d * vals[2]);

            if (fs[t - 1] < fs[t])
                t--;
        }
    }

    return (vals[0] + vals[3]) / 2;
}

/**
 * @brief Armijo-Goldstein 线性搜索
 *
 * @tparam T
 * @param[in] f
 * @param[in] df
 * @param[in] x 初始位置
 * @param[in] d 下降方向
 * @param[in] tol
 * @return double
 */
template <class T> inline double armijo_goldstein_line_search(double (*f)(T), T (*df)(T), T x, T d, double tol)
{
    double a = 0, b = std::numeric_limits<double>::infinity();
    double alpha = 10;
    double rho = 0.25;
    double t = 2;
    double fx = f(x);
    T dfx = df(x);

    while (b - a > tol)
    {
        double falpha = f(x + d * alpha);

        // 右侧判断，满足则更新左端；否则更新右端
        if (falpha < fx + alpha * rho * dot(dfx, d))
        {
            // 左侧判断，满足则直接退出
            if (falpha > fx + alpha * (1 - rho) * dot(dfx, d))
                break;

            // 否则增大左端
            a = alpha;

            // 如果右端无穷大，则可以增大试探步长。否则减小右端
            if (b == std::numeric_limits<double>::infinity())
            {
                alpha *= t;
                continue;
            }
        }
        else
            b = alpha;

        // 调整试探步长，使其靠近 [a, b] 的中点
        alpha = (a + b) / 2;
    }

    return alpha;
}

/**
 * @brief Wolfe-Powell 线性搜索
 *
 * @tparam T
 * @param[in] f
 * @param[in] df
 * @param[in] x
 * @param[in] d
 * @param[in] tol
 * @return double
 */
template <class T> inline double wolfe_powell_line_search(double (*f)(T), T (*df)(T), T x, T d, double tol)
{
    double a = 0, b = std::numeric_limits<double>::infinity();
    double alpha = 10;
    double rho = 0.25;
    double sigma = 0.5;
    double t = 2;
    double fx = f(x);
    T dfx = df(x);

    while (b - a > tol)
    {
        double falpha = f(x + d * alpha);
        T dfalpha = df(x + d * alpha);

        // 右侧判断，满足则更新左端；否则更新右端
        if (falpha < fx + alpha * rho * dot(dfx, d))
        {
            if (dot(dfalpha, d) > sigma * dot(dfx, d))
                break;

            // 否则增大左端
            a = alpha;

            // 如果右端无穷大，则可以增大试探步长。否则减小右端
            if (b == std::numeric_limits<double>::infinity())
            {
                alpha *= t;
                continue;
            }
        }
        else
            b = alpha;

        // 调整试探步长，使其靠近 [a, b] 的中点
        alpha = (a + b) / 2;
    }

    return alpha;
}

/**
 * @brief 回溯线性搜索
 *
 * @tparam T
 * @param[in] f
 * @param[in] df
 * @param[in] x
 * @param[in] d
 * @param[in] tol
 * @return double
 */
template <class T> inline double backtracking_line_search(double (*f)(T), T (*df)(T), T x, T d, double tol)
{
    double a = 0, b = std::numeric_limits<double>::infinity();
    double alpha = 10;
    double rho = 0.25;
    double sigma = 0.5;
    double t = 2;
    double fx = f(x);
    T dfx = df(x);

    while (b - a > tol)
    {
        double falpha = f(x + d * alpha);
        T dfalpha = df(x + d * alpha);

        // 右侧判断，满足则更新左端；否则更新右端
        if (falpha < fx + alpha * rho * dot(dfx, d))
            break;

        b = alpha;
        alpha = -dot(dfx, d) * alpha * alpha / (2 * (falpha - fx - dot(dfx, d) * alpha));
    }

    return alpha;
}

} // namespace math

} // namespace xi
